from pocketflow import AsyncFlow
from typing import Optional, Dict

from .microflow.ingest import IngestAction, UpsertAction, Completed
from .microflow.embedding import EmbeddingAction

from scorpio.core.common import get_logger

logger = get_logger(__name__)

class IngestionAgent:
    """文档摄取Agent，使用pocketflow AsyncFlow实现文档数据摄入"""
    
    shared: Dict
    _async_flow: Optional[AsyncFlow] = None
    
    def __init__(self, shared: Dict, max_iteration: int = 5):
        self.name = "ingestion"
        self.shared = shared
        
        # 初始化pocketflow节点
        self.ingest = IngestAction()
        self.embedding = EmbeddingAction()
        self.upsert = UpsertAction()
        self.completed = Completed()
    
    def create_flow(self) -> AsyncFlow:
        """创建文档摄取流程"""
        
        # 连接流程节点
        self.ingest - "failed" >> self.completed
        self.ingest - "prepared" >> self.embedding
        
        self.embedding - "embedded" >> self.upsert
        self.embedding - "failure" >> self.completed
        
        self.upsert - "success" >> self.completed
        self.upsert - "failure" >> self.completed
        
        # 创建并返回流程
        self._async_flow = AsyncFlow(start=self.ingest)
        return self._async_flow
    
    async def run(self):
        logger.info("start to run flow: IngestionFlow")
        return await self._async_flow.run_async(self.shared)
